module ModuleQTRANSITION

	use GLOBALS
	use ModuleQERRORS
	use ModuleSaving
	
contains

	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!!				Initialize Jacobian				!!!!!		
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	subroutine Compute_JACOB
	
		implicit none
        
        ! Local variable declarations
		real(8) :: rit_t(T_TR), mtit_t(T_TR), amt_t(T_TR), xmt_t(T_TR), ar_t(T_TR), xr_t(T_TR)	! variable transformations of r and mt
        real(8), dimension(T_TR) :: amtin_TR, arin_TR											! variable transformations of r and mt
		real(8) :: Jdstep																		! for marginal variable change
		integer :: tt, tt_aux																	! counter
		real(8), dimension(NS,T_TR) :: a_pol_TEMP, h_pol_TEMP									! auxiliary policy functions

		! variable transformations r and mt
		ar = (rub-r_LR)/(rub-rlb)
		xr = log(ar/(1.0d0-ar))
		amt = (mtub-mt_LR)/(mtub-mtlb)
		xmt = log(amt/(1.0D0-amt))
		
		! ED for m and r constant at steady state value for entire transition
		xr_t    = xr
		rit_t   = r_LR
		xmt_t   = xmt
		mtit_t  = mt_LR
	
		ED_ss =  ED_TR(xr_t,xmt_t)

		! Compute the Value and Policy Functions when the shock in r is last period
        Jdstep = 1e-5
		tt = T_TR

		! shock in r
		print *, 'Compute policy functions shock last period: r'
		xr_t      = xr
		xmt_t     = xmt
		xr_t(tt)  = xr*(1.0D0+Jdstep)
    
		arin_TR   = exp(xr_t)/(1.0d0+exp(xr_t))
		r_TR	  = arin_TR*rlb + (1.0d0-arin_TR)*rub
		amtin_TR  = exp(xmt_t)/(1.0D0+exp(xmt_t))
		mt_TR     = amtin_TR*mtlb + (1.0D0-amtin_TR)*mtub
		wge_TR    = (1.0D0-alpha)/(r_TR+delta)
		wge_TR    = alpha*(wge_TR**((1d0-alpha)/alpha))

		call Model_VF_TR
		a_pol_TEMP = a_pol_TR
		h_pol_TEMP = h_pol_TR

        ! Compute Jacobian in every period
		do tt = 1, T_TR
	
            ! Paths of transformed variables
			xr_t      = xr					! path for r: final steady-state
			xmt_t     = xmt					! path for m: final steady-state
			xr_t(tt)  = xr*(1.0D0+Jdstep)	! marginally change r in period tt
			
            ! Paths for true variables
			arin_TR   = exp(xr_t)/(1.0d0+exp(xr_t))
			r_TR	  = arin_TR*rlb + (1.0d0-arin_TR)*rub
			amtin_TR  = exp(xmt_t)/(1.0D0+exp(xmt_t))
			mt_TR     = amtin_TR*mtlb + (1.0D0-amtin_TR)*mtub
			wge_TR    = (1.0D0-alpha)/(r_TR+delta)
			wge_TR    = alpha*(wge_TR**((1d0-alpha)/alpha))

			! Allocate policy functions: a_pol, h_pol
			! Before the shock, use policies based on distance from the shock
			do tt_aux = 1, tt
				a_pol_TR(:,tt_aux) = a_pol_TEMP(:,tt_aux + T_TR - tt)
				h_pol_TR(:,tt_aux) = h_pol_TEMP(:,tt_aux + T_TR - tt)
			enddo

			! After the shock, steady-state policies
			do tt_aux = tt+1, T_TR
				a_pol_TR(:,tt_aux) = a_pol_lr
				h_pol_TR(:,tt_aux) = h_pol_lr
            enddo

            ! Aggregate 
			call ComputeMEASURE_TR
			call ComputeAGGREGATES_TR

            ! Eqm conditions
			ED_t(1:T_TR)          = asset_TR
			ED_t(T_TR+1:2*T_TR)   = BC_TR

            ! Change in eqm conditions relative to baseline to get Jacobian
			JACmat_t(:,tt) = (ED_t-ED_ss)/(Jdstep*xr)

        enddo

        ! Compute the Value and Policy Functions when the shock in mt is last period
		tt = T_TR

		! shock in mt
		print *, 'Compute policy functions shock last period: mt'
		xr_t      = xr
		xmt_t     = xmt
		xmt_t(tt) = xmt*(1.0D0+Jdstep)

		arin_TR   = exp(xr_t)/(1.0d0+exp(xr_t))
		r_TR	  = arin_TR*rlb + (1.0d0-arin_TR)*rub
		amtin_TR  = exp(xmt_t)/(1.0D0+exp(xmt_t))
		mt_TR     = amtin_TR*mtlb + (1.0D0-amtin_TR)*mtub
		wge_TR    = (1.0D0-alpha)/(r_TR+delta)
		wge_TR    = alpha*(wge_TR**((1d0-alpha)/alpha))

		call Model_VF_TR
		a_pol_TEMP = a_pol_TR
		h_pol_TEMP = h_pol_TR
    
        ! Compute Jacobian in every period
		do tt = 1, T_TR
    
			! Paths of transformed variables
			xr_t       = xr						! path for r: final steady-state
			xmt_t      = xmt					! path for mt: final steady-state
			xmt_t(tt)  = xmt*(1.0D0+Jdstep)		! marginally change mt in period tt

			! Paths for true variables
			arin_TR   = exp(xr_t)/(1.0d0+exp(xr_t))
			r_TR	  = arin_TR*rlb + (1.0d0-arin_TR)*rub
			amtin_TR  = exp(xmt_t)/(1.0D0+exp(xmt_t))
			mt_TR     = amtin_TR*mtlb + (1.0D0-amtin_TR)*mtub
			wge_TR    = (1.0D0-alpha)/(r_TR+delta)
			wge_TR    = alpha*(wge_TR**((1d0-alpha)/alpha))

			! Allocate policy functions: a_pol, h_pol
			! Before the shock, use policies based on distance from the shock
			do tt_aux = 1, tt
				a_pol_TR(:,tt_aux) = a_pol_TEMP(:,tt_aux + T_TR - tt)
				h_pol_TR(:,tt_aux) = h_pol_TEMP(:,tt_aux + T_TR - tt)
            enddo
            
			! After the shock, steady-state policies
			do tt_aux = tt+1, T_TR
				a_pol_TR(:,tt_aux) = a_pol_lr
				h_pol_TR(:,tt_aux) = h_pol_lr
            enddo

            ! Aggregate
			call ComputeMEASURE_TR
			call ComputeAGGREGATES_TR

            ! Eqm conditions
			ED_t(1:T_TR)          = asset_TR
			ED_t(T_TR+1:2*T_TR)   = BC_TR

            ! Change in eqm conditions relative to baseline to get Jacobian
			JACmat_t(:,T_TR+tt) = (ED_t-ED_ss)/(Jdstep*xmt)

		enddo
	
	end subroutine Compute_JACOB
	
		
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!! 				Quasi-Newton Step				!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	subroutine Compute_QNEWTON
	
		implicit none
        
        ! Local variable declarations
		real(8) :: rit_t(T_TR), mtit_t(T_TR), xmt_t(T_TR), xr_t(T_TR)
		real(8) :: xk(Kuvar), xkp1(Kuvar), xik(Kuvar), fk(Kuvar), fkp1(Kuvar), fik(Kuvar), Bjinvk(Kuvar,Kuvar), Bjinvkp1(Kuvar,Kuvar), eax(1), ek, ekp1, eik, eikm
		real(8) :: Btf(Kuvar), dx(Kuvar) 
		integer, parameter :: maxit = 30
		real(8) :: errf(1), wwx, err_t(maxit)	
		integer :: it, itbs, tt, itsave, maxitbs, itax, JACit
		
        ! prepare Newton
		tolf = 1e-5		! tolerance
        maxitbs = 10	! max backsteps
		wwx = 0.5D0		! slow updating
	
		! Transform r and mt to inputs for ED
		ar   = (rub-r_LR)/(rub-rlb)
		xr   = log(ar/(1.0d0-ar))
		amt = (mtub-mt_LR)/(mtub-mtlb)
		xmt = log(amt/(1.0D0-amt))
	
		! initialize in steady-state Jacobian
		xk(1:T_TR)          = xr					! initial guess: interest rate (transformed)
		xk(T_TR+1:2*T_TR)   = xmt					! initial guess: transfer level (transformed)
		Bjinvk              = inv_FET(JACmat_t)		! inverse of Jacobian
		fk                  = ED_ss					! error at initial guess
		eax                 = maxval(abs(fk));		! max absolute error at initial guess
		ek                  = eax(1)				! max absolute error at initial guess
	
        ! auxiliary variables for saving
		itsave = 10; JACit = 0

		do it = 1, maxit

			!!! new guess
			call matvec_FET(Bjinvk,fk,Btf)			! apply updating rule; matvec_FET(A,V,Y)
			xkp1 = xk - Btf							! apply updating rule
			xkp1 = wwx*xk + (1.0D0-wwx)*xkp1        ! slow updating
		
			dx   = xkp1 - xk						! difference between new and previous guess
		
			eikm = 99900000000.0D0					! initialize error from previous backstep

			!!! backstep check
			do itbs = 1, maxitbs
		
				xik   = xk + dx						! new guess

				xr_t   = xik(1:T_TR)				! separate guess for interest rate (transformed)
				xmt_t  = xik(T_TR+1:2*T_TR)			! separate guess for transfer level (transformed)
				fik    = ED_TR(xr_t,xmt_t)			! evaluate at new guess

				eax = maxval(abs(fik))				! max absolute error at new guess 
				eik = eax(1)						! max absolute error at new guess
			
				if ( eik < ek ) then 				! this was a good step, move forward
					exit
				elseif ( eikm < eik ) then			! last step was better, undo last shrinking step and exit
					dx  = 2.0D0*dx								
					exit
				endif
			
				! if you reached here, (eik > ek) but (eik < eikm) -> keep shrinking
				eikm = eik							! update error from backstepping
				dx   = dx/2.0D0 					! shrink			
				
            enddo  ! end loop in itbs
            
			xik = xk + dx							! updated guess

			xkp1 = xik; fkp1 = fik; ekp1 = eik;		! update guess, error, max absolute error
		
			errf   = maxval(abs(fkp1))				! max absolute error
			print *, '****************************'
			print *, '**  it = ', it, ', errf = ', errf
			print *, '****************************'
		
			if (errf(1) < 5e-4) then
				wwx = 0.5D0		! adjust slow updating weight
				maxitbs = 10	! adjust max number of backsteps
				!JACit   = 1		! adjust whether Jacobian is updated or not
			endif
		
			if (errf(1) < 6e-5) then
				wwx = 0.75D0	! adjust slow updating weight
				maxitbs = 20	! adjust max number of backsteps
				JACit   = 0		! adjust whether Jacobian is updated or not
			endif
		
			if (errf(1) < tolf ) then
				print *, 'equilibrium found'
                call SaveAggregates_TR
				exit
			elseif (errf(1) < 7e-5 .and. it > 1200) then	
				print *, 'exiting with errf(1) = ', errf(1)
				exit
			elseif (errf(1)>0.1d0 .and. it>20) then
				print *, 'Exit without saving at ', errf(1)
				exit
			else
				!!! update x and B
				if (JACit == 0) then
					Bjinvkp1 = Update_JACOBIAN(xk,xkp1,fk,fkp1,Bjinvk)			
				else
					Bjinvkp1 = Bjinvk 
				endif
			
				!!! rename stuff
				xk     = xkp1
				fk     = fkp1
				Bjinvk = Bjinvkp1
				ek     = ekp1
			endif	
		
		
		enddo  ! end loop in it 
		
	
	end subroutine Compute_QNEWTON
	
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!! 				  Update JACOBIAN				!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	function Update_JACOBIAN(xk,xkp1,fk,fkp1,Bk)  result(Bkp1)
	
		implicit none
        
        ! Local variable declarations
		real(8), intent(in) :: xk(Kuvar), xkp1(Kuvar), fk(Kuvar), fkp1(Kuvar), Bk(Kuvar,Kuvar)
		real(8)             :: Bkp1(Kuvar,Kuvar)
	
		real(8) :: dfk(Kuvar), ukvec(Kuvar), dkvec(Kuvar), Btf(Kuvar), Bnum1(Kuvar,1), Bnum2(1,Kuvar), Bnum(Kuvar,Kuvar), dkvecT(1,Kuvar), denB
	
		!!! update x and B
		dfk   = fkp1 - fk				! change in error
		call matvec_FET(Bk,dfk,ukvec)	! see Miranda Fackler page 38
		dkvec = xkp1 - xk				! change in guess

		!!! numerator computations
		Bnum1(:,1)  = dkvec - ukvec  ! Bnum1 = dk-uk -> a Nx1 "matrix"

		dkvecT(1,:) = dkvec			
		call matmul_FET(dkvecT,Bk,Bnum2) 	   ! this is Bnum2 = dk^T * Bk  -> a 1xN "matrix"   ~~ matmul_FET(A,B,C)			
		call matmul_FET(Bnum1,Bnum2,Bnum)      ! this is Bnum = (dk-uk)*(dk^T*Bk) -> a NxN matrix 			

		!!! denominator computations
		denB  = sum(dkvec*ukvec)

		!!! update Binv
		Bkp1 = Bk + (Bnum/denB)				
	
	end function Update_JACOBIAN
	
end module ModuleQTRANSITION